#models/detector.py
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.discriminant_analysis import StandardScaler
from torch.utils.data import Dataset


# 自定义Dataset类
class MyDataSet(Dataset):
    def __init__(self, csv_file):
        # 读取数据
        self.data = pd.read_csv(csv_file)
        
        # 提取特征（假设特征列为所有列，去掉标签列）
        self.features = self.data.drop(columns=['*label_Dos', '*label_Probe', '*label_R2L', '*label_U2R', '*label_normal'])
        self.features = self.features.values
        
        # 对特征进行标准化
        self.scaler = StandardScaler()
        self.features = self.scaler.fit_transform(self.features)

        # 提取标签
        self.labels = self.convert_labels(self.data)

    def __len__(self):
        # 返回数据集大小
        return len(self.data)

    def __getitem__(self, idx):
        # 获取每个样本的特征和标签
        x = torch.tensor(self.features[idx], dtype=torch.float32)
        y = torch.tensor(self.labels[idx], dtype=torch.float32)
        return x, y
    
    def convert_labels(self, df):
        # 将最后5列标签转化为数值类型
        labels = df[['*label_Dos', '*label_Probe', '*label_R2L', '*label_U2R', '*label_normal']].values
        return labels
class Detector(nn.Module):
    def __init__(self, input_dim):
        super(Detector, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(input_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1),
            nn.Sigmoid()  # 输出为正常流量的概率
        )

    def forward(self, x):
        return self.model(x)

class CNNDetector(nn.Module):
    def __init__(self, input_channels):
        super(CNNDetector, self).__init__()
        self.conv1 = nn.Conv1d(input_channels, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(64, 32, kernel_size=3, padding=1)
        self.fc1 = nn.Linear(32 * 10, 64)  # 假设输入大小为 (batch, channels, 10)
        self.fc2 = nn.Linear(64, 1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = torch.sigmoid(self.fc2(x))
        return x

class LSTMDetector(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers):
        super(LSTMDetector, self).__init__()
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, 1)

    def forward(self, x):
        _, (hn, _) = self.lstm(x)
        x = self.fc(hn[-1])
        return torch.sigmoid(x)
    
class GRUDetector(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layers):
        super(GRUDetector, self).__init__()
        self.gru = nn.GRU(input_dim, hidden_dim, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_dim, 1)

    def forward(self, x):
        _, hn = self.gru(x)
        x = self.fc(hn[-1])
        return torch.sigmoid(x)
        x = self.fc(hn[-1])
        return torch.sigmoid(x)
        return torch.sigmoid(x)
